from collections.abc import Iterator, Mapping
from typing import Any

from zkl_pyutils_fsspec import Fs, FsLike, resolve_child_fs, resolve_fs

from .loader import load_record, record_json_file_name


class RecordsScanner(Mapping[str, Any]):
    def __init__(self, fs: FsLike):
        self._fs = resolve_fs(fs)
        self._children_reader: dict[str, Any] | None = None

    def refresh(self, *, clean: bool = False):
        if clean or self._children_reader is None:
            self._children_reader = {}
            for path in scan_records_dir_path(self._fs):
                self._children_reader[path] = None
        else:
            for path in scan_records_dir_path(self._fs):
                if path not in self._children_reader:
                    self._children_reader[path] = None

    def __len__(self):
        if self._children_reader is None:
            self.refresh()
        return iter(self._children_reader.keys())

    def __iter__(self):
        if self._children_reader is None:
            self.refresh()
        return iter(self._children_reader.keys())

    def __getitem__(self, key: str):
        if self._children_reader is None:
            self.refresh()
        reader = self._children_reader[key]
        if reader is None:
            try:
                reader = load_record(resolve_child_fs(self._fs, key))
            except Exception:
                self._children_reader.pop(key)
            self._children_reader[key] = reader
        return reader


def scan_records_dir_path(fs: Fs, prefix: str = "") -> Iterator[str]:
    children_info = fs.ls(prefix)
    for info in children_info:
        if info['name'] == record_json_file_name:
            yield prefix
            return
    for info in children_info:
        if info['type'] == 'directory':
            yield from scan_records_dir_path(
                fs, f"{prefix}{fs.sep}{info['name']}")
